{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Assets hierarchy visualization\n",
    "\n",
    "\n",
    "## Quick links\n",
    "* Asset data dive notebook - [Hackathon github repo](https://github.com/cognitedata/open-industrial-data/blob/master/workshops/uni-hackathon/Part_3_Asset_Data_Dive.ipynb)\n",
    "* Documentation of [CDP concepts](https://doc.cognitedata.com/concepts/)\n",
    "* Reference documentation for the [Python SDK](https://cognite-sdk-python.readthedocs-hosted.com/en/latest/)\n",
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 0: Environment Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install the Cognite SDK package and some auxiliary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if you're working in google colab or similar\n",
    "!pip install -q cognite-sdk plotly networkx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Because of some limitations on google colab platform we should define huge auxiliary functions directly here, but we are working on eligant way for it\n",
    "\n",
    "By the way, if you know how to improve it - ideas or contributions are welcome :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load common.py\n",
    "\"\"\"\n",
    "\n",
    "DON'T BE SCARED\n",
    "\n",
    "You can don't dig into details here, just go to next cell\n",
    "\n",
    "\"\"\"\n",
    "from itertools import islice\n",
    "from typing import Dict, Iterable\n",
    "\n",
    "import networkx as nx\n",
    "import pandas as pd\n",
    "import plotly.graph_objs as go\n",
    "\n",
    "\n",
    "def sliding_window(seq: Iterable, n: int = 2):\n",
    "    \"\"\"\n",
    "\n",
    "    Generate an overlapping windows with variable size from iterator\n",
    "\n",
    "    Args:\n",
    "        seq:\n",
    "        n:\n",
    "\n",
    "    Returns:\n",
    "\n",
    "    \"\"\"\n",
    "    it = iter(seq)\n",
    "    result = tuple(islice(it, n))\n",
    "    if len(result) == n:\n",
    "        yield result\n",
    "    for elem in it:\n",
    "        result = result[1:] + (elem,)\n",
    "        yield result\n",
    "\n",
    "\n",
    "def make_assets_tree(df: pd.DataFrame) -> nx.DiGraph:\n",
    "    \"\"\"\n",
    "\n",
    "    Generates directional graph of assets from a given dataframe\n",
    "    \n",
    "    Args:\n",
    "        df:\n",
    "\n",
    "    Returns:\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    G = nx.DiGraph()\n",
    "    for path in df[\"path\"].values:\n",
    "        for parent_id, child_id in sliding_window(path, n=2):\n",
    "            G.add_edge(parent_id, child_id)\n",
    "\n",
    "    return G\n",
    "\n",
    "\n",
    "def hierarchy_pos(\n",
    "    graph: nx.DiGraph,\n",
    "    root: str = None,\n",
    "    width: float = 10000.0,\n",
    "    vert_gap: float = 0.2,\n",
    "    vert_loc: float = 0.0,\n",
    "    x_center=0.5,\n",
    "):\n",
    "    \"\"\"\n",
    "    If the graph is a tree this will return the positions to plot this in a\n",
    "    hierarchical layout.\n",
    "\n",
    "    G: the graph (must be a tree)\n",
    "\n",
    "    root: the root node of current branch\n",
    "    - if the tree is directed and this is not given, the root will be found and used\n",
    "    - if the tree is directed and this is given, then the positions will be just for the descendants of this node.\n",
    "    - if the tree is undirected and not given, then a random choice will be used.\n",
    "\n",
    "    width: horizontal space allocated for this branch - avoids overlap with other branches\n",
    "\n",
    "    vert_gap: gap between levels of hierarchy\n",
    "\n",
    "    vert_loc: vertical location of root\n",
    "\n",
    "    xcenter: horizontal location of root\n",
    "    \"\"\"\n",
    "    if not nx.is_tree(graph):\n",
    "        raise TypeError(\"cannot use hierarchy_pos on a graph that is not a tree\")\n",
    "\n",
    "    def _hierarchy_pos(\n",
    "        graph: nx.DiGraph,\n",
    "        root: str,\n",
    "        width: float = 1.0,\n",
    "        vert_gap: float = 0.2,\n",
    "        vert_loc: float = 0.0,\n",
    "        x_center: float = 0.5,\n",
    "        pos: Dict = None,\n",
    "        parent: str = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        see hierarchy_pos docstring for most arguments\n",
    "\n",
    "        pos: a dict saying where all nodes go if they have been assigned\n",
    "        parent: parent of this branch. - only affects it if non-directed\n",
    "\n",
    "        \"\"\"\n",
    "\n",
    "        if pos is None:\n",
    "            pos = {root: (x_center, vert_loc)}\n",
    "        else:\n",
    "            pos[root] = (x_center, vert_loc)\n",
    "        children = list(graph.neighbors(root))\n",
    "        if not isinstance(graph, nx.DiGraph) and parent is not None:\n",
    "            children.remove(parent)\n",
    "        if len(children) != 0:\n",
    "            dx = width / len(children)\n",
    "            next_x = x_center - width / 2 - dx / 2\n",
    "            for child in children:\n",
    "                next_x += dx\n",
    "                pos = _hierarchy_pos(\n",
    "                    graph,\n",
    "                    child,\n",
    "                    width=dx,\n",
    "                    vert_gap=vert_gap,\n",
    "                    vert_loc=vert_loc - vert_gap,\n",
    "                    x_center=next_x,\n",
    "                    pos=pos,\n",
    "                    parent=root,\n",
    "                )\n",
    "        return pos\n",
    "\n",
    "    return _hierarchy_pos(graph, root, width, vert_gap, vert_loc, x_center)\n",
    "\n",
    "\n",
    "def get_label(id_, client=None):\n",
    "    \"\"\" Get asset's name by given asset id \"\"\"\n",
    "    asset_info = client.assets.get_asset(id_)\n",
    "    return asset_info.to_json()[\"name\"]\n",
    "\n",
    "\n",
    "def make_assets_tree_plot(df: pd.DataFrame, root_id: int = None, max_depth: int = None) -> go.Figure:\n",
    "    \"\"\"\n",
    "\n",
    "    Generates assets plots from Assets dataframe\n",
    "\n",
    "    Args:\n",
    "        df:\n",
    "        root_id:\n",
    "        max_depth:\n",
    "\n",
    "    Returns:\n",
    "\n",
    "    \"\"\"\n",
    "    assets_tree = make_assets_tree(df)\n",
    "\n",
    "    if root_id is None:\n",
    "        root_id = next(iter(nx.topological_sort(assets_tree)))  # allows back compatibility with nx version 1.11\n",
    "\n",
    "    assets_tree = nx.dfs_tree(assets_tree, source=root_id, depth_limit=max_depth)\n",
    "    pos = hierarchy_pos(assets_tree, root=root_id)\n",
    "\n",
    "    # extract node coordinates and labels\n",
    "    Xn = [pos[i][0] for i in pos.keys()]\n",
    "    Yn = [pos[i][1] for i in pos.keys()]\n",
    "    labels = [get_label(id_) for id_ in pos.keys()]\n",
    "\n",
    "    # extract edges from tree\n",
    "    Xe = list()\n",
    "    Ye = list()\n",
    "    for e in assets_tree.edges():\n",
    "        Xe.extend([pos[e[0]][0], pos[e[1]][0], None])\n",
    "        Ye.extend([pos[e[0]][1], pos[e[1]][1], None])\n",
    "\n",
    "    # make plotly traces\n",
    "    trace_nodes = dict(\n",
    "        type=\"scatter\",\n",
    "        x=Xn,\n",
    "        y=Yn,\n",
    "        mode=\"markers\",\n",
    "        marker=dict(size=20, color=\"rgb(0, 0, 204)\"),\n",
    "        text=labels,\n",
    "        hoverinfo=\"text\",\n",
    "    )\n",
    "    trace_edges = dict(\n",
    "        type=\"scatter\", mode=\"lines\", x=Xe, y=Ye, line=dict(width=1, color=\"rgb(25,25,25)\"), hoverinfo=\"none\"\n",
    "    )\n",
    "\n",
    "    # some pretty details\n",
    "    axis = dict(\n",
    "        showline=False,  # hide axis line, grid, ticklabels and  title\n",
    "        zeroline=False,\n",
    "        showgrid=False,\n",
    "        showticklabels=False,\n",
    "        title=\"\",\n",
    "    )\n",
    "\n",
    "    layout = go.Layout(autosize=True, showlegend=False, xaxis=axis, yaxis=axis, hovermode=\"closest\")\n",
    "    fig = go.Figure(data=[trace_edges, trace_nodes], layout=layout)\n",
    "    return fig\n",
    "\n",
    "\n",
    "def configure_plotly_browser_state():\n",
    "    \"\"\"\n",
    "\n",
    "    Resolves an issue with plotly in google colab\n",
    "\n",
    "    Returns:\n",
    "\n",
    "    \"\"\"\n",
    "    from IPython.core.display import display, HTML\n",
    "\n",
    "    display(\n",
    "        HTML(\n",
    "            \"\"\"\n",
    "        <script src=\"/static/components/requirejs/require.js\"></script>\n",
    "        <script>\n",
    "          requirejs.config({\n",
    "            paths: {\n",
    "              base: '/static/base',\n",
    "              plotly: 'https://cdn.plot.ly/plotly-latest.min.js?noext',\n",
    "            },\n",
    "          });\n",
    "        </script>\n",
    "        \"\"\"\n",
    "        )\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Import the required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "from datetime import datetime, timedelta\n",
    "from datetime import datetime\n",
    "from getpass import getpass\n",
    "from typing import List, Any\n",
    "from itertools import islice\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "from cognite import CogniteClient\n",
    "\n",
    "import networkx as nx\n",
    "from networkx.algorithms.traversal.depth_first_search import dfs_tree\n",
    "import plotly.graph_objs as go\n",
    "from plotly.offline import iplot, init_notebook_mode\n",
    "init_notebook_mode(connected=True)\n",
    "\n",
    "\n",
    "pd.set_option('display.max_rows', 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Connect to the Cognite Data Platform\n",
    "The SDK client is the entrypoint to all data in CDP, and simply requires the API key that you generated in Part 1.\n",
    "\n",
    "When prompted for your API key, use the key generated by open industrial data as mentioned in the Getting Started steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = CogniteClient(api_key=getpass(\"Open Industrial Data API-KEY: \"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Learn about Organizing Industrial Data\n",
    "\n",
    "The Cognite Data Platform organizes digital information about the physical world. The building blocks of this representation are called *resources*, which you can read up on in detail [here](https://doc.cognitedata.com/concepts/#core-concepts).\n",
    "\n",
    "An important resource to understand is the Asset resource. This is the foundation for organizing industrial data -- time series, work orders, event logs and arbitrary files -- from across complex industrial systems.\n",
    "Assets are linked together with parent-child relationships to build a top-down hierarchical tree, known as \"The Asset Hierarchy\".\n",
    "For example, an Asset Hierarchy could look like this:\n",
    "```\n",
    "  Gas Export Compressor\n",
    "    |- First stage export compressor\n",
    "    |    |- Compressor\n",
    "    |    |- Scrubber\n",
    "    |    |- ...\n",
    "    |- Second stage export compressor\n",
    "    |- ...\n",
    "```\n",
    "Timeseries, events, files and other resources are attached to each Asset.\n",
    "\n",
    "The hierarchical structure can make it easier to find the timeseries data that you're looking for. Though there are [other ways](https://doc.cognitedata.com/concepts/#_3d-models-and-revisions) to do this, we'll focus on using the hierarchy today!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download a sample of assets up to a certain depth in the hierarchy\n",
    "df_sample_assets = client.assets.get_assets(limit=1000, depth=6).to_pandas().sort_values('depth')\n",
    "df_sample_assets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configure_plotly_browser_state()\n",
    "\n",
    "fig = make_assets_tree_plot(df_sample_assets)\n",
    "iplot(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may found this overcomplicated, but if you wanna explore your assets in interactive way - that's a good way. Also, you can play with `root_id` and `max_depth` arguments for `make_assets_tree_plot()` function, and build only a branch in details\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sample_assets = client.assets.get_assets(limit=2000, depth=15).to_pandas().sort_values('depth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configure_plotly_browser_state()\n",
    "\n",
    "fig = make_assets_tree_plot(df_sample_assets, root_id=4518112062673878, max_depth=4)\n",
    "iplot(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
